/**
 * @file aimbot.h
 * @brief 自瞄类头文件
 * @license Copyright© 2022 HITwh HERO-RoboMaster Group
 */

#ifndef AIMBOT_H
#define AIMBOT_H

#include <algorithm>
#include <map>
#include <opencv2/opencv.hpp>
#include <string>
#include <utility>
#include <vector>

#include "../armor/armor.h"
#include "../../../tool/include/util_func.h"
#include "interfaces/srv/armor.hpp"
#include "interfaces/msg/readpack.hpp"
#include "interfaces/msg/enemy.hpp"
#include "interfaces/msg/tobase.hpp"
// #include "interfaces/msg/autoaim_recorder.hpp"

#define TOP_NUM 30

/**
 * @brief 自瞄状态
 */
enum AimState
{
    /// 搜索装甲板状态
    SEARCH_STATE = 1,

    /// 追踪装甲板状态
    ARMOR_TRACK_STATE = 2,

    /// 追踪车辆状态
    CAR_TRACK_STATE = 3,

};

/**
 * @brief 装甲板检测类
 */
class AimBot
{
public:
    /// 陀螺速度阈值,超过该值之后自动进入反小陀螺模式
    constexpr static double GYRO_LOW_SPEED_THRESH = 0.5;

    ///
    constexpr static double GYRO_HIGH_SPEED_THRESH = 1.5;

public:
    int width = 1280;

    int height = 960; // 图片尺寸

    // 上一次按下右键的时间
    int last_click_right = 0;

    // 开始丢帧的时间
    int begin_lost_time = 0;

    /// 自瞄当前模式，在开始时执行赋值
    int curr_state;

    /// 下一帧自瞄模式,在结束时执行赋值
    int next_state;

    /// 记录装甲板丢失帧数
    int armor_lost_cnt = 0;

    /// 记录装甲板丢失后寻找统一辆车的帧数
    int car_lost_cnt = 0;

    /// 追踪状态下最大允许丢失帧数
    // TODO 最大允许丢失帧数需要测试
    int armor_max_lost_cnt = 6;
    int car_max_lost_cnt = 6;

    /// 搜寻成功标志
    bool targetSearched;

    /// 追踪正常进行中（最大丢帧数小于max_lost_cnt）
    bool if_next_frame_track = false;

    /// 当前追踪的数字
    int curr_number = 0;

    /// 是否使用数字识别
    int USE_MODEL;

    /// 是否使用ROI
    int ROI_ENABLE;

    /// 装甲板相对于相机绝对位置，单位m
    double absoluteDistance = 0.0;

    /// 候选灯条
    std::vector<Lightbar> lightbars;

    /// 候选装甲板, 每次在run函数运行开始时清空
    std::vector<Armor> armors;

    /// 当前是否找到装甲板
    bool armorFound;

    /// 当前识别到的待打击装甲板
    Armor target_armor;

    /// 当前识别到的待打击装甲板（副）
    Armor target_armor_deputy;

    /// 上一帧识别的装甲板
    Armor old_armor;

    /// ROI区域
    cv::Rect roi_rect;

    /// 历史坐标点
    CircleQueue<Target, TOP_NUM> his_targets;

    ///
    bool track_status;

    /// 存储最近5个装甲板大小的数据：1为大装甲板，0为小装甲板
    CircleQueue<int, 6> armor_type_recorder;

    bool is_balance_infantry = false;

    int big_armor_cnt = 0;

    /// 追踪车辆失败
    bool flag_track_car_failed = false;

    /// 切换装甲板
    bool flag_switch_armor;

    // // 进行第n个装甲板切换满足
    int switch_num = 0;

    // // 切换装甲板多次判断,满了之后才能判断成装甲板切换
    // bool flag_switch_armor_array[3]={false};

    // // 切换装甲板临时存储的yaw
    // double switch_armor_yaw[4]={0.0};

    // double switch_armor_time[4]={0.0};

    // //
    // std::vector<double> armor_yaws;
    // std::vector<double> armor_times;
#ifdef _USE_NETWORK
    /// Armor神经网络推理类的构建
    ArmorDetector armor_detector;
#endif

public:
    /**
     * @brief 默认构造函数
     */
    AimBot();

    /**
     * @brief 默认析构函数
     */
    ~AimBot();

    /**
     * @brief 初始化函数
     *
     * @param file_storage 参数配置文件
     */
    void init();

    /**
     * @brief 执行装甲板检测的核心函数
     *
     * @param src 源图像
     * @param mcu_data 通信包
     */
    void run(std::vector<Armor> armors, ReadPack work_msg, ArmorDir dir, const Pose &pose, double yaw_speed);

    /**
     * @brief 循环进入自瞄模式的一些设置
     * @details 比如根据上一帧的状态设置当前帧的状态等
     */
    void onEnter();

    /**
     * @brief 循环退出自瞄的一些操作
     * @details 比如清空一些数据容器
     */
    void onExit();

    /**
     * @brief 寻找装甲板模式，全图搜索
     */
    bool searchArmor(std::vector<Armor> armors, ReadPack work_msg);

    //    bool searchArmor_dynatic(cv::Mat &image, const ReadPack &mcu_data, const Pose &pose);

    /**
     * @brief 追踪装甲板模式，roi区域或者是全图检索
     */
    void trackArmor(std::vector<Armor> armors, ReadPack work_msg, ArmorDir dir, const Pose &pose, double yaw_speed);

    /**
     * @brief 追踪旋转前哨站
     */
    //    void trackArmor_dynamic(cv::Mat &image, const ReadPack &mcu_data, ArmorDir dir, const Pose &pose, double yaw_speed);

    /**
     * @brief 加载参数
     */
    void loadParam();

    /**
     * @brief 使用神经网路推理调用的函数
     */
    std::vector<Armor> &detectArmor(cv::Mat &image, int enemy_color, bool keep_gray_armor = false);

    // /**
    //  * @brief 根据对方装甲板颜色, 将图像预处理成二值图像
    //  *
    //  * @param src 源图像
    //  * @param mcu_color 敌方颜色
    //  */
    // void processImage(const cv::Mat &src, const int mcu_color);

    // /**
    //  * @brief 寻找可能灯条
    //  */
    // void findLightbars(const cv::Mat &processed_image, const int mcu_color);

    // /**
    //  * @brief 灯条匹配装甲板
    //  */
    // void matchArmors(const std::vector<Lightbar> &lightbars, std::vector<Armor> &armors, const Pose &pose);

    /**
     * @brief 装甲板打分
     */
    void gradeArmors(std::vector<Armor> &armors);

    /**
     * @brief 搜索模式下数字识别过滤和排序
     */
    bool trySearchArmor(std::vector<Armor> &armors, int mcu_number, int mcu_color);

    /**
     * @brief 利用位置追踪装甲板，在数字识别开启的情况下同时用数字识别过滤
     * @details 注意里面有反陀螺操作
     */
    void try_track(std::vector<Armor> &armors, ArmorDir dir, int color, double yaw_speed);

    /**
     * @brief 利用位置追踪装甲板，在数字识别开启的情况下同时用数字识别过滤
     * @details 注意里面有反陀螺操作 设计目的主要是将误识别和装甲板切换判断进行解耦
     * @date 2023.3.18
     * @author 井智（有问题欢迎交流）
     */
    void try_track_new(std::vector<Armor> &armors, ArmorDir dir, int color, double yaw_speed);

    /**
     * @brief 根据装甲板那的误差分数和优先级进行排序
     */
    void sortArmors(std::vector<Armor> &armors);

    /**
     * @brief 根据距离来决定是否使用数字识别
     * @details
     * 因为近距离会存在曝光过高光晕严重影响数字的情况，所以距离小于1.5m不使用数字识别
     * @details
     */
    void ApplyModelBaseOnDistance();

    /**
     * @brief 判断装甲板是否切换
     */
    bool judgesArmorSwitched(const Armor &curr_armor, const Armor &old_armor, ArmorDir dir, double yaw_speed);

    /**
     * @brief 在陀螺状态下判断装甲板切换 2023
     */
    int judgesArmorSwitched_antigyro(const Armor &curr_armor, const Armor &old_armor, double yaw_speed);

    /**
     * @brief 防止roi区域超出图像区域（640 * 480）
     */
    void preventROIExceed(int &x, int &y, int &width, int &height);

    /**
     * @brief 在图像上画出识别灯条
     */
    void drawBars(cv::Mat &canvas, std::vector<Lightbar> &bars);

    /**
     * @brief 在图像上画出装甲板
     * @details 待击打装甲板绿色，其余候选装甲板红色,
     * old_armor（要追踪的装甲板）白色
     */
    void drawArmors(cv::Mat &canvas, std::vector<Armor> &armors);

    /**
     * @brief 把一些识别信息放到显示到图片上
     * @details 信息例如：自瞄模式/是否使用数字识别/是否爆发式自瞄
     */
    void displayInfo(cv::Mat &canvas);

    /**
     * @brief 用于从其他模式转换到自瞄模式重新运行一些设置,例如工作状态的初始化
     */
    void config();

    const cv::Rect &getRoiRect() const;

    void setRoiRect(const cv::Rect &roiRect);

    void judge_armor_type(const CircleQueue<int, 6> &armor_type_recorder, bool &isBigArmor);

    /**
     * @brief 执行旋转前哨站装甲板检测的核心函数
     *
     * @param src 源图像
     * @param mcu_data 通信包
     */
    void run_dynamic(std::vector<Armor> armors, ReadPack work_msg, ArmorDir dir, const Pose &pose, double yaw_speed);

    bool searchArmor_dynatic(std::vector<Armor> armors, ReadPack work_msg);

    /**
     * @brief 追踪旋转前哨站
     */
    void trackArmor_dynamic(std::vector<Armor> armors, ReadPack work_msg, ArmorDir dir, const Pose &pose, double yaw_speed);

    void try_track_dynamic(std::vector<Armor> &armors, ArmorDir dir, int color, double yaw_speed);

    bool judgesArmorSwitchedDynamic(const Armor &curr_armor, const Armor &old_armor, ArmorDir dir, double yaw_speed);

public:
    /// 原图像的浅拷贝，用于debug
    cv::Mat canvas;

    /// 预处理得到的灰度图像, 用于寻找灯条
    cv::Mat processed_image;

    /// ROI图片
    cv::Mat roi_image;

    /// 灰度图
    cv::Mat gray_image;

    /// 二值图
    cv::Mat binary_image;

    /// 通道相减图
    cv::Mat subtract_image;

    /// 颜色通道二值图
    cv::Mat channel_thresh_image;

    /// 轮廓被视为灯条的最大小面积
    double MAX_LIGHTBAR_AREA;
    double MIN_LIGHTBAR_AREA;

    /// 灯条长短比下限
    double MIN_LIGHTBAR_RATIO;

    /// 大装甲板长短比下限
    double MIN_ASPECT_BIG_RATIO;

    /// 大装甲板长短比上限
    double MAX_ASPECT_BIG_RATIO;

    /// 小装甲板长短比下限
    double MIN_ASPECT_SMALL_RATIO;

    /// 小装甲板长短比上限
    double MAX_ASPECT_SMALL_RATIO;

    /// 两灯条长度比上限
    double MAX_LENGTH_RATIO;

    /// 灯条倾斜度上限
    double MAX_LIGHTBAR_ANGLE;

    /// 灯条倾斜度之差上限
    double MAX_LIGHTBAR_DELTA;

    /// 装甲板倾斜度上限
    double MAX_ARMOR_ANGLE;

    /// 每帧送进armor.cpp和分类器的装甲板的最大数量
    int MAX_CANDIDATE_NUM;

    /// 源图像二值化阈值
    int BLUE_BINARY_THRES;
    int RED_BINARY_THRES;

    /// 颜色通道阈值
    int BLUE_CHANNEL_THRESH;
    int RED_CHANNEL_THRESH;

    /// 通道相减图像灰度阈值,用于第一种图像预处理方法
    int BLUE_SUBTRACT_THRES;
    int RED_SUBTRACT_THRES;

    /// 算子核大小
    int KERNEL_SIZE;

    /// 相机的fy，由内参矩阵得出
    double camera_fy;

#ifdef DISTORTION_CORRECT
    Mat camera_internal_matrix;
    Mat distortion_coeff;
#endif // DISTORTION_CORRECT

#ifndef COMPILE_WITH_CUDA
    cv::Mat kernel;
#else
    cv::Ptr<cv::cuda::Filter> kernel;
#endif // COMPILE_WITH_CUDA
};

extern double getYawResolution(const Target &target, const bool isBigArmor);
extern double getPitchResolution(const Target &target);

#endif /// AIMBOT_H
